package com.apex.flink.sink;

import com.apex.flink.FlinkEnvironment;
import com.apex.flink.stream.FlinkStreamSink;
import com.apex.flink.utils.ConfigKeyName;
import com.typesafe.config.Config;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.typeutils.RowTypeInfo;
import org.apache.flink.connectors.kudu.connector.KuduTableInfo;
import org.apache.flink.connectors.kudu.connector.writer.AbstractSingleOperationMapper;
import org.apache.flink.connectors.kudu.connector.writer.KuduWriterConfig;
import org.apache.flink.connectors.kudu.connector.writer.RowOperationMapper;
import org.apache.flink.connectors.kudu.streaming.KuduSink;
import org.apache.flink.connectors.kudu.table.KuduCatalog;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSink;
import org.apache.flink.types.Row;
import org.apache.kudu.ColumnSchema;
import org.apache.kudu.Type;
import org.apache.kudu.client.CreateTableOptions;

import java.util.*;

/**
 * kudu 数据写入方法
 */
public class KuduStreamSink implements FlinkStreamSink<Row,Row> {
    private Config config;
    private String tableName;
    private String KUDU_MASTERS;
    private List<String> primaryKey = new ArrayList<>();
    private List<String> hashKey = new ArrayList<>();
    private int bucket;
    private int replicas;

    /**
     * 数据写入方法
     * @param env env
     * @param dataStream dataStream
     * @return DataStreamSink<Row>
     */
    @Override
    public DataStreamSink<Row> outputStream(FlinkEnvironment env, DataStream<Row> dataStream) {
        //获取字段类型
        RowTypeInfo rowTypeInfo = (RowTypeInfo)dataStream.getType();
        //初始化kudu配置
        KuduWriterConfig writerConfig = KuduWriterConfig.Builder.setMasters(KUDU_MASTERS).build();
        KuduCatalog catalog = new KuduCatalog(KUDU_MASTERS);
        //获取列信息
        List<String> columns = new ArrayList<>();
        columns.addAll(Arrays.asList(rowTypeInfo.getFieldNames()));
        //初始化kudu表信息
        KuduTableInfo tableInfo = KuduTableInfo.forTable(tableName);

        tableInfo.createTableIfNotExists(
                () ->{
                   String [] fields = rowTypeInfo.getFieldNames();
                   TypeInformation [] types = rowTypeInfo.getFieldTypes();
                   List<ColumnSchema> list = new ArrayList<>();
                   for (int i = 0 ; i < fields.length ; i ++){
                    switch (types[i].toString().toUpperCase()){
                        case "STRING":
                            if (primaryKey.contains(fields[i])){
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.STRING).key(true).build());
                            }else {
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.STRING).build());
                            }
                            break;
                        case "BOOLEAN":
                            list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.BOOL).build());
                            break;
                        case "TINYINT":
                            if (primaryKey.contains(fields[i])){
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.INT8).key(true).build());
                            }else {
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.INT8).build());
                            }
                            break;
                        case "SMALLINT":
                            if (primaryKey.contains(fields[i])){
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.INT16).key(true).build());
                            }else {
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.INT16).build());
                            }
                            break;
                        case "INT":
                            if (primaryKey.contains(fields[i])){
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.INT32).key(true).build());
                            }else {
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.INT32).build());
                            }
                            break;
                        case "BIGINT":
                            if (primaryKey.contains(fields[i])){
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.INT64).key(true).build());
                            }else {
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.INT64).build());
                            }
                            break;
                        case "FLOAT":
                            if (primaryKey.contains(fields[i])){
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.FLOAT).key(true).build());
                            }else {
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.FLOAT).build());
                            }
                            break;
                        case "DOUBLE":
                            if (primaryKey.contains(fields[i])){
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.DOUBLE).key(true).build());
                            }else {
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.DOUBLE).build());
                            }
                            break;
                        case "BYTES":
                            list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.BINARY).build());
                            break;
                        case "TIMESTAMP(3)":
                            if (primaryKey.contains(fields[i])){
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.UNIXTIME_MICROS).key(true).build());
                            }else {
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.UNIXTIME_MICROS).build());
                            }
                            break;
                        case "DECIMAL":
                            if (primaryKey.contains(fields[i])){
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.DECIMAL).key(true).build());
                            }else {
                                list.add(new ColumnSchema.ColumnSchemaBuilder(fields[i],Type.DECIMAL).build());
                            }
                            break;
                            default:
                                System.out.println("no support type ....");
                                System.exit(0);
                    }

                   }
                   System.out.println("find field is size : "+list.size());
                   return list;
                },
                () -> {
                    return new CreateTableOptions()
                            .setNumReplicas(replicas)
                            .addHashPartitions(hashKey, bucket);
                });
        try {
            catalog.createTable(tableInfo,true);
        }catch (Exception e){
            e.printStackTrace();
            System.exit(0);
        }

        //创建写入方式
        RowOperationMapper rowOperationMapper = new RowOperationMapper(columns.toArray(new String[0]),
                AbstractSingleOperationMapper.KuduOperation.UPSERT);
        //写入数据到kudu
        return dataStream.addSink(new KuduSink<Row>(writerConfig,tableInfo,rowOperationMapper));
    }

    /**
     *
     * @param environment
     */
    @Override
    public void prepare(FlinkEnvironment environment) {
        String keys = "";
        String hashKeys = "";
        if (config.hasPath(ConfigKeyName.SINK_KUDU_MASTER)){
            KUDU_MASTERS = config.getString(ConfigKeyName.SINK_KUDU_MASTER);
            tableName = config.getString(ConfigKeyName.SINK_KUDU_TABLE_NAME);
            bucket = config.getInt(ConfigKeyName.SINK_KUDU_TABLE_BUCKET);
            replicas = config.getInt(ConfigKeyName.SINK_KUDU_TABLE_REPLICAS);
            keys = config.getString(ConfigKeyName.SINK_KUDU_TABLE_PRIMARYKEY);
            hashKeys = config.getString(ConfigKeyName.SINK_KUDU_TABLE_HASHKEY);
        }
        primaryKey.addAll(Arrays.asList(keys.split(",")));
        hashKey.addAll(Arrays.asList(hashKeys.split(",")));
    }

    /**
     *
     * @return
     */
    @Override
    public Config getConfig() {
        return config;
    }

    /**
     *
     * @param config config
     */
    @Override
    public void setConfig(Config config) {
        this.config = config;
    }
}
